import numpy as np

from opt_func.FuncSet import *


class Function:
    def __init__(self, function, values_range, data_index=None, restrains=None, restrains_all_rebase_score=1e8):
        """
        优化函数对象的初始化定义，包括：
        1.优化函数
        2.各变量取值空间
        3.变量次序（用于多树训练）
        4.约束列表（元素都是约束函数）
        5.违反约束时的惩罚数值
        """
        function_list = ['Rastrigin', 'MishraBird',
                         'GoldsteinPrice', 'Rosenbrock']
        if isinstance(function, str):
            if function not in function_list:
                raise AttributeError('输入的函数不存在')
            else:
                self.function = eval(function)
        else:
            self.function = function
        if data_index is None:
            raise ValueError("请制定变量序列")
        self.restrains_all_rebase_score = restrains_all_rebase_score
        self.values_range = values_range
        self.data_index = data_index
        self.restrains = restrains

    def __call__(self, fixed, simulation=None, **kwargs):
        """:
        将对象作为函数进行调用
        在该方法中，实现了：
        1.生成优化函数传入数组（这里的实现不怎么优美，可以看看怎么优化）
        2.当有约束时考虑约束的影响
        """
        # TODO:对于违反程度做个判定方法
        # TODO:如果所有都违反约束怎么办
        fixed = np.array(fixed)
        data = fixed
        # 如果进行了蒙特卡洛仿真，则对数组进行堆叠构造
        if simulation is not None:
            _, wide = simulation.shape
            fixed = fixed.repeat(wide).reshape(-1, wide)
            data = np.vstack([fixed, simulation])
        data = data[self.data_index]
        score = np.array(self.function(data, **kwargs))
        # 约束的处理
        score = self.process_restrain(data, score, **kwargs)
        return [score.min(), score.mean()]

    def process_restrain(self, data, score, **kwargs):
        if isinstance(self.restrains, dict):  # 以字典形式反映惩罚系数
            is_rebase_restrains = np.array([restrain(data, **kwargs)
                                            for restrain in self.restrains.keys()])*np.array(
                self.restrains.values())
            score += is_rebase_restrains.sum(axis=0)
        elif isinstance(self.restrains, list):  # 没有惩罚系数
            is_rebase_restrains = np.array([restrain(data, **kwargs)
                                            for restrain in self.restrains]).any(axis=0)
            score = score[is_rebase_restrains]
        # TODO:决定全部违反时怎么办
            if is_rebase_restrains.sum() == 0:
                score = np.array([self.restrains_all_rebase_score])
        return score
